﻿namespace CatClient.Message.Spi.Heartbeat.Extend.Netstat
{
    using CatClient.Util;
    using System;
    using System.Collections.Generic;
    using System.Linq;
    using System.Net;
    using System.Runtime.InteropServices;

    public class NetstatManager
    {
        private const int AF_INET = 2;

        public List<TcpProcessRecord> GetAllTcpConnections()
        {
            int pdwSize = 0;
            List<TcpProcessRecord> source = new List<TcpProcessRecord>();
            GetExtendedTcpTable(IntPtr.Zero, ref pdwSize, true, 2, TcpTableClass.TCP_TABLE_OWNER_PID_ALL, 0);
            IntPtr pTcpTable = Marshal.AllocHGlobal(pdwSize);
            try
            {
                if (GetExtendedTcpTable(pTcpTable, ref pdwSize, true, 2, TcpTableClass.TCP_TABLE_OWNER_PID_ALL, 0) != 0)
                {
                    return new List<TcpProcessRecord>();
                }
                MIB_TCPTABLE_OWNER_PID mib_tcptable_owner_pid = (MIB_TCPTABLE_OWNER_PID) Marshal.PtrToStructure(pTcpTable, typeof(MIB_TCPTABLE_OWNER_PID));
                IntPtr ptr = (IntPtr) (((long) pTcpTable) + Marshal.SizeOf(mib_tcptable_owner_pid.dwNumEntries));
                for (int i = 0; i < mib_tcptable_owner_pid.dwNumEntries; i++)
                {
                    MIB_TCPROW_OWNER_PID structure = (MIB_TCPROW_OWNER_PID) Marshal.PtrToStructure(ptr, typeof(MIB_TCPROW_OWNER_PID));
                    byte[] buffer1 = new byte[] { structure.localPort[1], structure.localPort[0] };
                    byte[] buffer2 = new byte[] { structure.remotePort[1], structure.remotePort[0] };
                    source.Add(new TcpProcessRecord(new IPAddress((long) structure.localAddr), new IPAddress((long) structure.remoteAddr), BitConverter.ToUInt16(buffer1, 0), BitConverter.ToUInt16(buffer2, 0), structure.owningPid, structure.state));
                    ptr = (IntPtr) (((long) ptr) + Marshal.SizeOf(structure));
                }
            }
            catch (OutOfMemoryException exception)
            {
                LogInfoWriter.GetInstance("Cat", 0, LogInfoWriter.LogLayout.Default).Error("内存溢出，获取TCP连接信息失败：" + exception.Message, exception);
            }
            catch (Exception exception2)
            {
                LogInfoWriter.GetInstance("Cat", 0, LogInfoWriter.LogLayout.Default).Error("获取TCP连接信息失败：" + exception2.Message, exception2);
            }
            finally
            {
                Marshal.FreeHGlobal(pTcpTable);
            }
            return source.Distinct<TcpProcessRecord>().ToList<TcpProcessRecord>();
        }

        public List<UdpProcessRecord> GetAllUdpConnections()
        {
            int pdwSize = 0;
            List<UdpProcessRecord> source = new List<UdpProcessRecord>();
            GetExtendedUdpTable(IntPtr.Zero, ref pdwSize, true, 2, UdpTableClass.UDP_TABLE_OWNER_PID, 0);
            IntPtr pUdpTable = Marshal.AllocHGlobal(pdwSize);
            try
            {
                if (GetExtendedUdpTable(pUdpTable, ref pdwSize, true, 2, UdpTableClass.UDP_TABLE_OWNER_PID, 0) != 0)
                {
                    return new List<UdpProcessRecord>();
                }
                MIB_UDPTABLE_OWNER_PID mib_udptable_owner_pid = (MIB_UDPTABLE_OWNER_PID) Marshal.PtrToStructure(pUdpTable, typeof(MIB_UDPTABLE_OWNER_PID));
                IntPtr ptr = (IntPtr) (((long) pUdpTable) + Marshal.SizeOf(mib_udptable_owner_pid.dwNumEntries));
                for (int i = 0; i < mib_udptable_owner_pid.dwNumEntries; i++)
                {
                    MIB_UDPROW_OWNER_PID structure = (MIB_UDPROW_OWNER_PID) Marshal.PtrToStructure(ptr, typeof(MIB_UDPROW_OWNER_PID));
                    byte[] buffer1 = new byte[] { structure.localPort[1], structure.localPort[0] };
                    source.Add(new UdpProcessRecord(new IPAddress((long) structure.localAddr), BitConverter.ToUInt16(buffer1, 0), structure.owningPid));
                    ptr = (IntPtr) (((long) ptr) + Marshal.SizeOf(structure));
                }
            }
            catch (OutOfMemoryException exception)
            {
                LogInfoWriter.GetInstance("Cat", 0, LogInfoWriter.LogLayout.Default).Error("内存溢出，获取UDP连接信息失败：" + exception.Message, exception);
            }
            catch (Exception exception2)
            {
                LogInfoWriter.GetInstance("Cat", 0, LogInfoWriter.LogLayout.Default).Error("获取UDP连接信息失败：" + exception2.Message, exception2);
            }
            finally
            {
                Marshal.FreeHGlobal(pUdpTable);
            }
            return source.Distinct<UdpProcessRecord>().ToList<UdpProcessRecord>();
        }

        [DllImport("iphlpapi.dll", CharSet=CharSet.Auto, SetLastError=true)]
        private static extern uint GetExtendedTcpTable(IntPtr pTcpTable, ref int pdwSize, bool bOrder, int ulAf, TcpTableClass tableClass, uint reserved = 0);
        [DllImport("iphlpapi.dll", CharSet=CharSet.Auto, SetLastError=true)]
        private static extern uint GetExtendedUdpTable(IntPtr pUdpTable, ref int pdwSize, bool bOrder, int ulAf, UdpTableClass tableClass, uint reserved = 0);
    }
}

